﻿//#define VERBOSE

using System;
using RayDen.Library.Core;
using RayDen.Library.Core.Primitives;
using RayDen.Library.Entity.Scene;
using RayDen.RayEngine.Core;
using RayDen.RayEngine.Core.Devices;
using RayDen.RayEngine.Core.Interface;
using RayDen.RayEngine.Core.Threads;
using RayDen.RayEngine.Core.Types;
using RayDen.RayEngine.Managers;
using RayDen.RayEngine.Samplers;

namespace RayDen.RayEngine.Engines.PathTracer {
    // ReSharper disable RedundantThisQualifier
    
    public static class SampleDataOffset {
        public static int BsdfSampleOffset = 0;
        public static int SingleLightSampleOffset = 1;
        public static int LightSampleOffset1 = 2;
        public static int LightSampleOffset2 = 3;
        public static int LightSampleOffset3 = 4;


        public static int BrdfSampleOffset1 = 5;
        public static int BrdfSampleOffset2 = 6;
        public static int BrdfSampleOffset3 = 7;

        public static int RussianRouletteOffset = 8;


        public static int SamplesPerVertex = 9;
    }

    public class MCPathTracer : BaseRenderingEngine {
        public MCPathTracer(RayEngineScene scene) : base(scene) { }

        public override RenderThread CreateRenderThread(int threadIndex, IntersectionEngine intersectionDevice, RayEngineScene scene, IImageFilm film, IFrameManager frameManager, SamplingContext samplingContext)
        {
            return new DeviceRenderThread<PathBuffer<MCPathTracerPath>>(threadIndex, intersectionDevice, scene,
                                                                        LowLatencySampling, film, samplingContext,frameManager);
        }

        public override ISampler CreateSampler(int width, int height, int maxPathDepth, int threadIndex, int screenStartLine, int windowHeight) {
            var sampler = new MutatorSampler() { lowLatency = LowLatencySampling, sceneMaxPathDepth = maxPathDepth };
//            var sampler = new RandomSampler() { lowLatency = LowLatencySampling, tIdx = threadIndex, screenStartLine = screenStartLine };

            return sampler;
        }

        public override Core.Interface.IImageFilm CreateFilm(int width, int height) {
            return new MCImageFilm(width, height);
        }
    }


    public class MCImageFilm : ImageFilm<MCPathTracerPath> {
        public float AcceptanceWeight = 0.25f;
        public MCImageFilm(int width, int height) : base(width, height) { }

        public override void Consume(PathSamplerBase pt) {
            var path = (MCPathTracerPath)pt;
            this.Samples.SplatSample(path.Sample, path.Radiance, path.pathWeight);
            // Check if the sample buffer is full
            if (this.Samples.IsFull()) {
                this.statsTotalSampleCount += this.Samples.GetSampleCount();

                // Splat all samples on the film
                this.SplatSampleBuffer(this.IsPreviewOver(), this.Samples);
                this.Samples = this.GetFreeSampleBuffer();
            }

            if (path.pathWeight < 1f & (path.pathWeight > AcceptanceWeight | path.mutateCandidate) & !path.storedSample) {
                path.storedSample = true;
                path.mutateCandidate = false;
                path.mutationsCount = 0;
            }

        }
    }

    public class MCPathTracerPath : PathSamplerBase {
        internal RgbSpectrum Throughput;
        internal bool mutateCandidate;
        public PathTracerPathState PathState;



        protected int depth, tracedShadowRayCount;
        public float pathWeight;
        public bool storedSample;
        public int mutationsCount;
        protected bool specularBounce;

        protected RayEngineScene scene;

        public int MaxRaysPerPath {
            get;
            set;
        }

        struct ShadowRayInfo {
            public float pdf;
            public RgbSpectrum color;
            public RayData shadowRay;
            public int currentShadowRayIndex;
        }
        private ShadowRayInfo[] secRays;


        public override void InitPath(IPathProcessor buffer) {
            base.InitPath(buffer);
            this.scene = pathIntegrator.Scene;

            if (storedSample && this.Sample != null) {
                this.mutationsCount++;
                var sampler = (MutatorSampler)pathIntegrator.Sampler;
                this.Sample = sampler.Mutate(this.Sample, mutationsCount);
            }
            else 
            {
                this.storedSample = false;
                this.mutationsCount = 0;
                this.mutateCandidate = false;
                this.Sample = pathIntegrator.Sampler.GetSample(null);
                //this.Sample.InitSample(scene.MaxPathDepth * SampleDataOffset.SamplesPerVertex);
            }

            this.Radiance = new RgbSpectrum(0f);
            this.Throughput = new RgbSpectrum(1f);
            this.PathState = PathTracerPathState.EyeVertex;
            this.secRays = new ShadowRayInfo[scene.ShadowRayCount];
            IRay ray;
            pathIntegrator.Scene.Camera.GetRay(Sample.imageX, Sample.imageY, out ray);
            this.PathRay = (RayData)ray;
            this.RayIndex = -1;
            this.pathWeight = 1.0f;
            this.tracedShadowRayCount = 0;
            this.depth = 0;
            this.specularBounce = true;
        }

        public override bool FillRayBuffer(RayBuffer rayBuffer) {
            var leftSpace = rayBuffer.LeftSpace();
            if (((PathState == PathTracerPathState.EyeVertex) && (1 > leftSpace)) ||
            ((PathState == PathTracerPathState.ShadowRaysOnly) && (tracedShadowRayCount > leftSpace)) ||
            ((PathState == PathTracerPathState.NextVertex) && (tracedShadowRayCount + 1 > leftSpace)))
                return false;
            if (PathState != PathTracerPathState.ShadowRaysOnly)
                RayIndex = rayBuffer.AddRay(PathRay);
            if (PathState == PathTracerPathState.NextVertex || PathState == PathTracerPathState.ShadowRaysOnly) {
                for (int i = 0; i < tracedShadowRayCount; ++i)
                    secRays[i].currentShadowRayIndex = rayBuffer.AddRay(secRays[i].shadowRay);
            }
            return true;
        }


        public override void Advance(RayBuffer rayBuffer, SampleBuffer consumer) {
            int currentTriangleIndex = 0;
            SurfaceIntersectionData hitInfo = null;
#if VERBOSE

            try {
#endif

            base.Advance(rayBuffer, consumer);

            RayHit rayHit = rayBuffer.rayHits[RayIndex];

            if (((PathState == PathTracerPathState.ShadowRaysOnly) || (PathState == PathTracerPathState.NextVertex)) &&
                (tracedShadowRayCount > 0)) {
                for (int i = 0; i < tracedShadowRayCount; ++i) {
                    RayHit shadowRayHit = rayBuffer.rayHits[secRays[i].currentShadowRayIndex];
                    RgbSpectrum attenuation;
                    if (this.ShadowRayTest(ref shadowRayHit, out attenuation)) {
                        //                            Radiance.MADD()
                        Radiance += attenuation * ((secRays[i].color) / secRays[i].pdf);
                        pathWeight *= 1f/secRays[i].pdf;
                    }
                }
                Splat(consumer);
                
                return;
            }

            depth++;
            bool missed = rayHit.Index == 0xffffffffu;
            if (missed || PathState == PathTracerPathState.ShadowRaysOnly || depth > scene.MaxPathDepth) {
                if (specularBounce && missed) {
                    Radiance += this.SampleEnvironment(PathRay.Dir) * Throughput;
                }
                Splat(consumer);
                
                return;
            }

            // Something was hit
            hitInfo = SurfaceSampler.GetIntersection(ref PathRay, ref rayHit);
            currentTriangleIndex = (int)rayHit.Index;

            //If Hit light)
            if (hitInfo.IsLight) {
                if (specularBounce) {
                    var lt = scene.GetLightByIndex(currentTriangleIndex);
                    var le = (RgbSpectrum)(lt.Le(ref PathRay.Dir));
                    Radiance += Throughput * le;
                }
                Splat(consumer);
                
                return;
            }

            var hitPoint = PathRay.Point(rayHit.Distance);

            Vector wo = -PathRay.Dir;
            tracedShadowRayCount = 0;
            if (hitInfo.MMaterial.IsDiffuse()) {
                float lightStrategyPdf = scene.ShadowRayCount / (float)scene.Lights.Length;
                RgbSpectrum lightTroughtput = Throughput * hitInfo.Color;
                for (int i = 0; i < scene.ShadowRayCount; ++i) {
                    int currentLightIndex = scene.SampleLights(Sample.GetLazyValue((depth) * SampleDataOffset.SingleLightSampleOffset));
                    var light = scene.Lights[currentLightIndex];

                    var ls = new LightSample();
                    light.EvaluateShadow(ref hitPoint, ref hitInfo.ShadingNormal,
                        Sample.GetLazyValue((depth) * SampleDataOffset.LightSampleOffset1),
                        Sample.GetLazyValue((depth) * SampleDataOffset.LightSampleOffset2),
                        Sample.GetLazyValue((depth) * SampleDataOffset.LightSampleOffset3), ref ls);
                    if (ls.Pdf <= 0f) {
                        continue;
                    }

                    secRays[tracedShadowRayCount].color = (RgbSpectrum)(ls.Spectrum);
                    secRays[tracedShadowRayCount].pdf = ls.Pdf;
                    secRays[tracedShadowRayCount].shadowRay = ls.LightRay;

                    Vector lwi = secRays[tracedShadowRayCount].shadowRay.Dir;
                    RgbSpectrum fs;
                    hitInfo.MMaterial.f(
                        ref secRays[tracedShadowRayCount].shadowRay.Dir,
                        ref wo, ref hitInfo.ShadingNormal, ref Throughput, out fs);
                    secRays[tracedShadowRayCount].color *= lightTroughtput*
                                                           Vector.AbsDot(ref hitInfo.ShadingNormal, ref lwi)*
                                                           fs;
                    if (!secRays[tracedShadowRayCount].color.IsBlack()) {
                        secRays[tracedShadowRayCount].pdf *= lightStrategyPdf;
                        tracedShadowRayCount++;
                    }
                }
            }

            float fPdf;
            Vector wi;

            RgbSpectrum f = hitInfo.MMaterial.Sample_f(ref wo, out wi, ref hitInfo.Normal, ref hitInfo.ShadingNormal, ref Throughput,
                                                      Sample.GetLazyValue((depth) * SampleDataOffset.BrdfSampleOffset1), Sample.GetLazyValue((depth) * SampleDataOffset.BrdfSampleOffset2),
                                                      Sample.GetLazyValue((depth) * SampleDataOffset.BrdfSampleOffset3), ref hitInfo.TextureData, 
                                                      out fPdf, out specularBounce);

            if ((fPdf <= 0.0f) || f.IsBlack()) {
                if (tracedShadowRayCount > 0)
                    PathState = PathTracerPathState.ShadowRaysOnly;
                else {
                    Splat(consumer);
                    
                }
                return;
            }
            pathWeight *= fPdf;
            Throughput *= f / fPdf;

            if (depth > scene.MaxPathDepth) {
                float prob = Math.Max(Throughput.Filter(), scene.RussianRuletteImportanceCap);
                if (prob >= Sample.GetLazyValue((depth) * SampleDataOffset.RussianRouletteOffset)) {
                    Throughput /= prob;
                    pathWeight *= prob;
                }
                else {
                    if (tracedShadowRayCount > 0)
                        PathState = PathTracerPathState.ShadowRaysOnly;
                    else {
                        Splat(consumer);
                        
                    }

                    return;
                }
            }

            PathRay.Org = hitPoint;
            PathRay.Dir = wi.Normalize();
            PathState = PathTracerPathState.NextVertex;
            mutateCandidate = !storedSample && pathWeight > 0.1f && tracedShadowRayCount > 0;

#if VERBOSE
            }
            catch (Exception ex) {
                RayDen.Library.Components.SystemComponents.Tracer.Trace(currentTriangleIndex.ToString() + "\r\n", 2);
                RayDen.Library.Components.SystemComponents.Tracer.TraceLine(ex.Message);
            }
#endif
        }

        private bool ShadowRayTest(ref RayHit shadowRayHit, out RgbSpectrum attenuation) {
            var hit = shadowRayHit.Index == 0xffffffffu || scene.IsLight((int)shadowRayHit.Index);
            attenuation = new RgbSpectrum(1f);
            if (hit) return true;
            var mat = scene.MatLib.GetSurfMat(scene.GetMeshByTriangleIndex((int)shadowRayHit.Index).MaterialName);

            if (mat.AlphaTexture != null)
            {
                attenuation = (RgbSpectrum)(mat.AlphaTexture as ConstTextureInfo).Color;
                hit = attenuation.Filter() < Sample.GetLazyValue();
            }

            return hit;
        }

        private RgbSpectrum SampleEnvironment(Vector vector) {
            return this.scene.SampleEnvironment(-vector);
        }
    }
    // ReSharper restore RedundantThisQualifier
}